Text Generation Inference 簡稱 TGI，是由 Hugging Face 開發的 LLM Inference 框架。其中整合了相當多推論技術，例如 Flash Attention, Paged Attention, Continuous Batching 以及 BNB & GPTQ Quantization 等等，加上 Hugging Face 團隊強大的開發能量與活躍的社群參與，使 TGI 成為部署 LLM Service 的最佳選擇之一。
今天就來介紹 TGI 的用法吧！
TGI 提供 Local Install 跟 Docker Image 兩種用法，但是 Local Install 要花上非常久的時間編譯與安裝，而且也很容易遇到環境配置的問題，除非要客製化停用某些套件之類的，不然一律推薦使用 Docker Image 操作。
可以從官方 GitHub Packages 查看 TGI Docker Image 有哪些版本，或者直接使用 latest 最新版本：
docker pull ghcr.io/huggingface/text-generation-inference:latest
筆者目前使用的是 sha-5ba53d4 版，因為這個 Docker Image 大小約 10GB 左右，所以下載需要花一段時間。完成下載之後可以透過 --help 查看參數說明：
docker run --rm ghcr.io/huggingface/text-generation-inference --help
我們可以先用個小模型 OPT-125M 簡單測試一下：
docker run --gpus all --shm-size 1g \
-p 8080:80 -v $PWD/data:/data \
ghcr.io/huggingface/text-generation-inference \
--model-id facebook/opt-125m
等到訊息紀錄出現 Connected 就代表服務啟動完成，我們可以透過 curl 測試：
curl -X POST 127.0.0.1:8080/generate \
-d '{"inputs":"Hello, "}' \
-H 'Content-Type: application/json'
得到類似以下的輸出：
{"generated_text":" I'm a new player and I'm looking for a good team to play with. I'm a"}
TGI 會透過網路下載模型，並且把模型轉換成 Safetensors 格式，然後將模型權重存在容器內部的 /data 路徑裡面。可以使用參數 -v $PWD/data:/data 將此路徑 Mapping 出來，這樣就不用每次執行的時候都要重新下載一遍模型權重。反過來說，我們也可以先將模型下載好，然後再 Mapping 進容器裡面，例如：
# launch_tgi.sh
MODEL_PATH=Models/opt-125m
git clone https://huggingface.co/facebook/opt-125m $MODEL_PATH
docker run --gpus all --shm-size 1g \
-p 8080:80 -v $PWD/$MODEL_PATH:/$MODEL_PATH \
ghcr.io/huggingface/text-generation-inference \
--model-id /$MODEL_PATH
TGI 的 API 用法可以參考說明文件，這裡介紹幾個比較重要的參數：
將 decoder_input_details 和 details 設為 True 會回傳一些詳細資訊，例如生成 Token 總數、輸入 Prompt 的 Token 總數等等。給 stop 一個 List of String 可以控制輸出停止點。
最特別的是 TGI 可以設定 truncate 參數，讓系統幫你把太長的 Prompt 截斷，例如設定為 500 的話，就會把 Prompt 前面超過 500 個 Tokens 的內容全部切除，只剩最後面 500 Tokens 當輸入。
這在限制輸入長度很好用，例如做 Retrieval-Based Few-Shot Prompting 時，可以取非常多結果出來，把相似度較低的擺在前面，藉由 truncate 參數自然切掉相似度太低的範例，這樣就能在維持輸入長度不會超出系統限制的同時，盡可能的加上更多範例在 Few-Shot Prompt 裡面。
最後透過 Python 呼叫模型生成的程式碼大致如下：
import json
import requests
url = "http://localhost:8080/generate"
params = {
"inputs": "This is a long prompt maybe, ",
"parameters": {
"best_of": 1,
"details": True,
"return_full_text": True,
"decoder_input_details": True,
"truncate": 4, # 只保留最後四個 Tokens
"max_new_tokens": 128,
"stop": ["\n", "."],
"do_sample": True,
"temperature": 0.5,
"top_k": 10,
"top_p": 0.95,
},
}
resp = requests.post(url, json=params)
result = json.loads(resp.text)
print(result)
因為 return_full_text 被設為 True 所以回傳結果會包含 Prompt 與 Generation，雖然結果裡面看起來 Prompt 並沒有被截斷，但是仔細觀察 details 裡面的 prefill 就可以看到實際的 Prompt 其實是從 "maybe" 開始的。
在眾多 TGI 參數裡面，最值得注目的就是我們多了很多量化模型的選擇，透過 --quantize <QUANTIZE> 來指定要使用哪一種，目前 TGI 支援 AWQ, EETQ, BNB, GPTQ 等，以下簡單介紹這些參數選項：
bitsandbytes 8-Bit 量化，雖然速度偏慢，但還是支援最廣泛、穩定的選擇。bitsandbytes-nf4 4-Bit 量化，大部分的模型都可以直接使用此選項，資料型態為 BNB-NF4，可能更符合模型權重分佈的一種資料型態。bitsandbytes-fp4 4-Bit 量化，與 BNB-NF4 類似，但使用標準的 4-Bit 浮點數資料型態。gptq 4-Bit 量化，需要使用做過 GPTQ Post Training 的模型，可以到 HF Hub 上搜尋，例如 TheBloke 提供的 GPTQ 模型。awq 4-Bit 量化，類似 GPTQ 需要提供指定格式的模型，也可以參考 TheBloke 提供的 AWQ 模型。eetq 8-Bit 量化，應該可以直接用，但這個選項滿新的，感覺還有些 Bug 的樣子。雖然官方說準備棄用 BNB，但現階段來說 BNB 還是相對穩定一些。有了量化，我們單顯卡平民就天下無敵了！先拿個 Taiwan Llama 13B 熱熱身：
docker run --gpus all --shm-size 1g \
-p 8080:80 -v $PWD/data:/data \
ghcr.io/huggingface/text-generation-inference \
--model-id yentinglin/Taiwan-LLaMa-v1.0 \
--quantize bitsandbytes
順利跑起來之後，簡單拿個中文測試：
curl -X POST 127.0.0.1:8080/generate \
-d '{"inputs":"### USER: 嗨 ### ASSISTANT: "}' \
-H 'Content-Type: application/json'
# {"generated_text":"你好！我今天可以如何協助你？"}
好欸，看到了我們熟悉的母語！
因為 TGI 參數很多很複雜，於是筆者將常用的參數整理成一份 Python Script 放在此 GitHub Gist 上方便使用，此腳本僅供參考，請以自身運行的環境與 TGI 版本為主。
除了 Decoder 以外，像是 T5 或 M2M 之類的 Encoder-Decoder 也可以用透過 TGI 來運行。但如果不是 TGI 特別支援的模型，可能會回到原本 HF Transformers 的實做。一些像是 Tensor-Parallel 或 Flash Attention 的功能就無法用到，但依然有 Continuous Batching 或 Streaming Outputs 這些功能可以用。因此若要使用 Seq2Seq 模型也是可以考慮用 TGI 來部署。
若要更精細的部署模型，與 Token 相關的參數至關重要。對於小顯卡而言，可以減少記憶體消耗，對於大顯卡而言，則能充分利用記憶體空間。最主要的三個參數如下：
--max-input-length 單筆最大輸入長度。--max-total-tokens 單筆最大總長度。--max-batch-prefill-tokens 所有 Batch 加起來的最大輸入長度。其中 --max-batch-prefill-tokens 是影響記憶體消耗最關鍵的參數，會影響整個 Prefill 階段能夠容納多少 Token 當輸入。
在 Decoder-Only LM 裡面，通常會使用 Prefill 代表一開始的輸入，而 Generate 則代表後續 Autoregressive Decoding 的過程。一般而言 Transformers 會在 Prefill 階段消耗掉大量記憶體，而後續 Generate 的記憶體消耗增長速度則相對較緩。也就是說如果一次輸入一整篇長文到 LLM 裡面，可能會直接發生記憶體不足的錯誤。但如果是讓 LLM 慢慢生成一篇長文，則可以生成非常長的文章。
決定這些參數值的方向大概分成輸入長度、輸出長度與批次大小。在不同應用情境下，這些設定都會不太一樣。以 Llama 2 13B 為例，模型是以 4K Context Window 進行訓練的，那我們可以分配 3K 給輸入，剩下 1K 給輸出，因此得到：
--max-input-length 為 3000。--max-total-tokens 為 3000 + 1000 = 4000。註：一般而言 Context Length 的 2K, 4K 通常是以 1024 為單位，所以 3K/1K 的分配實際上會是
3 * 1024 = 3072跟1024個 Tokens。
若我們預計同時推論 4 筆輸入，則可以設定參數：
--max-batch-prefill-tokens 為 3000 * 4 = 12000
預設 --max-concurrent-requests 為 128 筆，在這樣的設定下，如果使用者每筆輸入都是 3000 個 Tokens，那 TGI 最多只會同時推論 4 個請求，剩下請求的都會被 Queue 起來，等待任何一個生成結束之後再拿出來放進 Batch 裡面繼續生成。
另外可以將 --max-best-of 設為 1 就好，因為我們通常只需要生成一筆結果。最後使用 bitsandbytes-nf4 做 4-Bit 量化，整個指令看起來會像這樣：
docker run --gpus all --shm-size 1g \
-p 8080:80 -v $PWD/data:/data \
ghcr.io/huggingface/text-generation-inference:latest \
--model-id TheBloke/Llama-2-13B-Chat-fp16 \
--quantize bitsandbytes-nf4 \
--max-best-of 1 \
--max-concurrent-requests 128 \
--max-input-length 3000 \
--max-total-tokens 4000 \
--max-batch-prefill-tokens 12000
不過 TGI 其實是動態決定當下的 Batch Size，例如使用者每筆輸入都減為 1500 個 Tokens 的話，那 Batch Size 可能就會增加到同時推論 8 筆。因此只要在不會超出記憶體的情況下，參數 --max-batch-prefill-tokens 能開多大就開多大，確保服務能夠同時處理的 Batch Size 為最大，便能提昇整個生成的吞吐量。
至於具體到底要調到多少才不會 OOM，那就只能慢慢測試了。筆者習慣以 2K 為單位往上遞增做測試，如果 8K 不會爆就試 10K，依此類推。
TGI Client 是用來跟 TGI Server 交流的客戶端介面，因為在 TGI 的文件沒有太多介紹，所以相當容易被人遺忘。首先透過 pip 安裝 text-generation 套件：
pip install text-generation
基本使用方法如下：
from text_generation import Client
client = Client("http://127.0.0.1:8080", timeout=600)
resp = client.generate(
"Hello",
max_new_tokens=16,
stop_sequences=["."],
do_sample=False,
truncate=2048,
)
print(resp.generated_text)
其實就是將 requests 的用法做個包裝，並且幫 Response 物件定義了明確的類別，因此輸入 resp. 的時候，會跳出 generated_text 的提示，開發的時候會更方便一點。透過 TGI Client 進行串流輸出也比較方便：
resp = client.generate_stream(
"Hello",
max_new_tokens=16,
stop_sequences=["."],
do_sample=False,
truncate=2048,
)
for chunk in resp:
print(end=chunk.generated_text, flush=True)
print()
更多詳細的用法，可以參考 PyPI 的介紹頁面。
使用空 Prompt 生成 128 Tokens，來比較一下 HF Transformers, vLLM 與 TGI 的速度：
HF FP16 - 1.23 ms
vLLM FP16 - 0.48 ms
TGI FP16 - 0.45 ms
TGI 的速度與 vLLM 差不多，但 TGI 的優勢在於量化支援較廣泛，因此來比較一下各種量化方式的速度：
FP16 - 0.45 ms
BNB 8-Bit - 0.64 ms
BNB 4-Bit - 0.73 ms
GPTQ 4-Bit - 0.80 ms
AWQ 4-Bit - 0.55 ms
有量化的模型雖然會略慢一些，但基本上都還是比 HF Transformers 快。這裡筆者沒有測試 EETQ 因為跑起來感覺不太正常，等這個量化方式穩定一點再來測試看看。
最後比較各框架單筆推論的速度：
HF FP16 - 27 ms
ggml FP16 - 20 ms
TGI FP16 - 18 ms
vLLM FP16 - 18 ms
其實只看一筆的話，大家的速度都還是差不多的。因此如果是那種放在個人電腦裡面，一人一個 LLM 的用途，ggml 的輕量化部署依然是個滿理想的選擇。
筆者曾經使用 TGI + Taiwan Llama 替人翻譯過一篇約 13000 Tokens 的英文文章，以下分享這個應用的實做。
首先我先將文章存在 content.txt 裡面，並手動分段，約兩三行的內容就放兩個換行。因為文章內容沒有很長，所以手動分段還算可以，主要是為了確保邊界正確。但如果應用在更長的文章上，可能需要考慮一些自動尋找 Boundary 的工具協助了。
接下來用 TGI 將 Taiwan Llama 架起來，使用 BNB-NF4 量化，接著撰寫以下程式碼進行翻譯：
import json
from concurrent.futures import ThreadPoolExecutor
import requests
from tqdm import tqdm
def main():
# 讀取文章並以 "\n\n" 切成多個 Chunks
with open("content.txt", "rt", encoding="UTF-8") as fp:
text = fp.read().strip()
text = text.split("\n\n")
# Taiwan Llama 提供的 Prompt Template
template = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: 請將以下句子從英文翻譯成中文: {} ASSISTANT:"
# 定義每段 Chunk 的翻譯函式
def translate(source):
prompt = template.format(source)
target = generate(prompt)
return {"Original": source, "Translate": target}
# 最多同時發送 128 個 Requests
results = list()
with ThreadPoolExecutor(max_workers=128) as executor:
with tqdm(total=len(text), ncols=80) as progress:
for res in executor.map(translate, text):
results.append(res)
progress.update()
# 將結果存成 JSON 檔
with open("results.json", "wt", encoding="UTF-8") as fp:
json.dump(results, fp, ensure_ascii=False, indent=4)
# 定義發送 HTTP Request 的函式
def generate(prompt):
url = "http://localhost:8080/"
# 參考 Taiwan Llama Demo 網頁的預設參數
data = {
"inputs": prompt,
"parameters": {
"do_sample": True,
"best_of": 1,
"max_new_tokens": 1000,
"stop": ["\n\n"],
"temperature": 0.7,
"top_k": 50,
"top_p": 0.90,
},
}
res = requests.post(url, json=data)
return json.loads(res.text)[0]["generated_text"]
if __name__ == "__main__":
main()
使用 RTX 3090 約四分半可以完成整份翻譯，給各位參考看看。
今天介紹了 TGI 推論框架，是許多推論服務採用的框架，他的速度與廣泛模型支援、量化技術等，使其成為無論是開發者還是研究者的首選框架。
其實除了這幾天介紹的 ONNX, ggml, vLLM, TGI 以外，還有非常多的推論框架，例如 Exllama V2, CTranslate 2 等等，這些推論框架都有其穩定的開發與用戶。
不過這些框架的變化速度也相當快，很有可能筆者這幾天介紹的框架、參數和用法等等，隔天一個大改就全部不能用了。（好沒有技術保存價值的划水文章 Q_Q
接下來會介紹 Offloading 的概念，明天見！
iThome鐵人賽
看影片追技術